# -*- coding: utf-8 -*-
# ----------------------------
# @Time    : 2021/9/23 上午11:15
# @Author  : acedar
# @FileName: lr_regression.py
# ----------------------------

# **** 评估指标：
# ********** rmse: root mean squared error
# ********** mse: mean squared error
# ********** r2: R2 metric
# ********** mae: mean absolute error
# ***********

from pyspark.sql import SparkSession
from pyspark.ml.regression import LinearRegression
from pyspark.ml.evaluation import RegressionEvaluator


spark = SparkSession.builder.appName("test").getOrCreate()
data_path = "../datasets/mllib/sample_linear_regression_data.txt"
data_df = spark.read.format("libsvm").load(data_path)
data_df.show(2)

train_df, test_df = data_df.randomSplit([0.8, 0.2])

lr = LinearRegression().setMaxIter(100).setRegParam(0.3).setElasticNetParam(0.8)
lrm = lr.fit(train_df)

trs = lrm.summary
print(f"train mse: {trs.meanSquaredError}")
print(f"train rmse: {trs.rootMeanSquaredError}")
print(f"train mae: {trs.meanAbsoluteError}")
print(f"train r2: {trs.r2}")

pred_df = lrm.transform(test_df)

evaluator = RegressionEvaluator().setLabelCol("label")\
    .setPredictionCol("prediction").setMetricName("mse")
acc = evaluator.evaluate(pred_df)
print(f"test mse: {acc}")


